'''
Esteganosaurio

Autores:
- Antonio Martin
- Rodrigo Olmo
'''

from random import randint

try:
    import Image
except:
    print "Se necesita la libreria PIL"
    exit()

BASE_AM = 8

def byteToPixel(pixel, byte):
    '''Recibe un pixel y un byte, devuelve el pixel modificado para contener
       la informacion de ese byte'''
    valor_r = pixel[0] + (byte & 0b00000111) / 0b00000001
    valor_g = pixel[1] + (byte & 0b00011000) / 0b00001000
    valor_b = pixel[2] + (byte & 0b11100000) / 0b00100000
    return (valor_r, valor_g, valor_b)

def pixelToByte(pixel):
    '''Recibe un pixel, devuelve el byte contenido en ese pixel'''
    return (pixel[0] % BASE_AM) * 0b00000001 \
         + (pixel[1] % BASE_AM) * 0b00001000 \
         + (pixel[2] % BASE_AM) * 0b00100000

def numToPixel(num):
    valor_r = (num & 0x0000FF) / 0x000001
    valor_g = (num & 0x00FF00) / 0x000100
    valor_b = (num & 0xFF0000) / 0x010000
    return (valor_r, valor_g, valor_b)

def pixelToNum(pixel):
    return pixel[0] * 0x000001 \
         + pixel[1] * 0x000100 \
         + pixel[2] * 0x010000
         
         
def codificar(f_entrada, f_cifrar, f_salida):
    imagen = Image.open( f_entrada ).convert("RGB")
    info = open( f_cifrar ,"rb")
    salida = imagen.point(lambda i: (i / BASE_AM) * BASE_AM)
    longitud = 0
    pixels = salida.load() # Con esto accedemos a los pixels mucho mas rapido
    for x in range(salida.size[0]):
        for y in range(salida.size[1]):
            if x == 0 and y == 0:
                continue
            leido = info.read(1)
            pixel = pixels[x,y]
            if len(leido) == 0: 
                pixels[x,y] = byteToPixel(pixel, randint(0, (2 ** BASE_AM) - 1))
            else:
                byte = ord(leido)
                nuevo_pixel = byteToPixel(pixel, byte)
                pixels[x,y] = nuevo_pixel
                longitud += 1
    if len(info.read(1)) != 0:
        print "Archivo a ocultar demasiado grande"
        return
    pixel_longitud = numToPixel(longitud)
    pixels[0,0] = pixel_longitud
    salida.save( f_salida )

def decodificar(f_entrada, f_salida):
    imagen = Image.open( f_entrada )
    info = open( f_salida , "wb")
    pixels = imagen.load()
    pixel_size = pixels[0,0]
    longitud = pixelToNum(pixel_size)
    for x in range(imagen.size[0]):
        for y in range(imagen.size[1]):
            if x == 0 and y == 0:
                continue
            if longitud == 0 : 
                break
            longitud-=1
            pixel = pixels[x,y]
            byte = pixelToByte(pixel)
            info.write(chr(byte))

if __name__ == "__main__":
    from sys import argv
    try:
        if "-c" == argv[1]:
            codificar(argv[2], argv[3], argv[4])
        elif "-d" == argv[1]:
            decodificar(argv[2], argv[3])
    except IndexError:
        print "Uso: " + argv[0] + " -c " + "<imagen_contenedora> " + "<archivo_a_cifrar> " + "<archivo_de_salida>"
        print "     " + argv[0] + " -d " + "<imagen_contenedora> " + "<archivo_a_crear>"
    
    
